"""Reads the blastclustoutput and divides proteins into given number of groups.

The proteins are assigned to a group randomly and each group is of roughly
equal size. This allows for performing a k-fold cross validation.
Usage:
    python group_blastclustoutput.py ../results/blastoutput 5
"""
import sys
import getopt
import random

def process(blastclustoutput, nr_groups):
    """Divide proteins into roughly equal groups and print to screen"""
    fh = open(blastclustoutput, 'r')

    clusters = fh.readlines()
    fh.close()
    nr_clusters = len(clusters)
    group_size = nr_clusters / nr_groups
    group_rest = nr_clusters % nr_groups

    groups = []
    for i in range(nr_groups):
        groups.append([])

    for i in range(len(groups)):
        if group_rest > 0:
            size = group_size + 1
            group_rest -= 1
        else:
            size = group_size
        
        for j in range(size):
            ri = random.randint(0, len(clusters) - 1)
            groups[i].append(clusters.pop(ri).split()[0])

    for i in range(len(groups)):
        print "Group" + str(i)

        for j in range(len(groups[i])):
            print groups[i][j]

def main():
    # parse command line options
    try:
        opts, args = getopt.getopt(sys.argv[1:], "h", ["help"])
    except getopt.error, msg:
        print msg
        print "for help use --help"
        sys.exit(2)
    # process options
    for o, a in opts:
        if o in ("-h", "--help"):
            print __doc__
            sys.exit(0)
    # process arguments
    if (args == 2):
        process(args[0], int(args[1])) # process() is defined elsewhere
    else:
        print __doc__

if __name__ == "__main__":
    main()
